/// Source : https://leetcode.com/problems/minimize-malware-spread-ii/
/// Author : liuyubobobo
/// Time   : 2018-10-21

#include <iostream>
#include <vector>
#include <unordered_set>
#include <unordered_map>

using namespace std;


/// Brute Force
/// Time Complexity: O(len(initial) * v^2)
/// Space Complexity: O(v)
class Solution {

private:
    int n;

public:
    int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {

        n = graph.size();
        vector<bool> removed(n, false);
        for(int v: initial)
            removed[v] = true;

        vector<unordered_set<int>> infection(n);
        for(int v: initial){
            vector<bool> visited(n, false);
            dfs(graph, v, removed, visited);
            for(int u = 0; u < n; u ++)
                if(visited[u])
                    infection[u].insert(v);
        }

        unordered_map<int, int> effective;
        for(int i = 0; i < n; i ++)
            if(infection[i].size() == 1)
                effective[*infection[i].begin()] ++;

        int res = -1, best = 0;
        for(const pair<int, int>& p: effective)
            if(p.second > best){
                best = p.second;
                res = p.first;
            }
            else if(p.second == best)
                res = min(res, p.first);
        return res;
    }

private:
    void dfs(const vector<vector<int>>& g, int v,
             const vector<bool>& removed, vector<bool>& visited){

        visited[v] = true;

        for(int i = 0; i < n; i ++)
            if(g[v][i] && !removed[i] && !visited[i])
                dfs(g, i, removed, visited);
    }
};


int main() {

    vector<vector<int>> g1 = {{1,1,0,0},{1,1,1,0},{0,1,1,1},{0,0,1,1}};
    vector<int> initial1 = {3,0};
    cout << Solution().minMalwareSpread(g1, initial1) << endl;
    // 0

    vector<vector<int>> g2 = {
            {1,0,0,0,0,0,0,0,0},
            {0,1,0,0,0,0,0,0,0},
            {0,0,1,0,1,0,1,0,0},
            {0,0,0,1,0,0,0,0,0},
            {0,0,1,0,1,0,0,0,0},
            {0,0,0,0,0,1,0,0,0},
            {0,0,1,0,0,0,1,0,0},
            {0,0,0,0,0,0,0,1,0},
            {0,0,0,0,0,0,0,0,1}};
    vector<int> initial2 = {6,0,4};
    cout << Solution().minMalwareSpread(g2, initial2) << endl;
    // 0

    vector<vector<int>> g3 = {{1,1,0},{1,1,1},{0,1,1}};
    vector<int> initial3 = {0,1};
    cout << Solution().minMalwareSpread(g3, initial3) << endl;
    // 1

    return 0;
}